# -*- coding: utf-8 -*-
"""
@author: administrator
"""
import shutil
from pathlib import Path

import torch.nn.functional
from torch.utils.data import DataLoader
from torchvision import transforms
from torchvision.datasets import CIFAR10
from torchvision.transforms import ToTensor

import utils.pytorch_utils as utils
from vgg16_net import Vgg16

if __name__ == '__main__':
    transform = transforms.Compose(
        [
            transforms.Resize(225),
            transforms.CenterCrop(224),
            transforms.ToTensor(),
            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
        ]
    )

    data_train = CIFAR10(root="data", train=True, transform=transform, download=True)
    data_test = CIFAR10(root="data", train=False, transform=transform, download=True)

    dataloader_train = DataLoader(data_train, batch_size=128, shuffle=True)
    dataloader_test = DataLoader(data_test, batch_size=128)

    device = utils.get_device()
    print("device:", device)
    net = Vgg16()
    net.to(device)
    net.train()

    # optimizer = optim.SGD(net.parameters(), lr=0.01, momentum=0.9, weight_decay=0.005)
    optimizer = torch.optim.Adam(net.classifier[6].parameters(), lr=0.001)
    shutil.rmtree("out")
    Path("out").mkdir(exist_ok=True)
    num_epochs = 10
    for i in range(num_epochs):
        running_loss = 0.0
        for index, (data, label) in enumerate(dataloader_train):
            data, label = data.to(device), label.to(device)
            optimizer.zero_grad()
            out = net(data)
            loss = torch.nn.functional.cross_entropy(out, label)
            loss.backward()
            optimizer.step()
            running_loss += loss.item()

        utils.save(net, "vgg16", i)
        print(f"Epoch {i + 1}/{num_epochs} Loss: {running_loss / len(dataloader_train)}")